## Load packages and source auxiliary functions
library(brms)
library(dplyr)
library(rio)
library(ggplot2)
library(gridExtra)
library(psych)
library(tidyr)
library(stevemisc)
library(sf)
library(urbnmapr)
library(RColorBrewer)
library(countrycode)
library(ivreg)
library(multiwayvcov)


fd <- function(model, variable, data, level=0.9, delta = c("sd", "range", "unit"), 
               dist = c("nb", "normal")){
  delta <- match.arg(delta)
  dist <- match.arg(dist)
  b <- as_draws_df(model)
  b_inds <- grep("b_", colnames(b), value=TRUE)
  if(dist == "nb"){
    shape_inds <- grep("b_shape_", colnames(b))
    G <- as.matrix(b[,shape_inds])
    Z <- as.matrix(cbind(1, data[,gsub("b_shape_", "", colnames(G)[-1])]))
  }
  if(dist == "normal"){
    shape_inds <- grep("b_sigma_", colnames(b))
    G <- as.matrix(b[,shape_inds])
    Z <- as.matrix(cbind(1, data[,gsub("b_sigma_", "", colnames(G)[-1])]))
  }
  mean_inds <- setdiff(b_inds, b_inds[shape_inds])
  B <- as.matrix(cbind(b[, mean_inds], 1))
  X <- as.matrix(cbind(1, data[,gsub("b_", "", 
                                     colnames(B)[-c(1, ncol(B))])], log(data$age)))
  X1 <- X2 <- X
  Z1 <- Z2 <- Z
  if(delta == "range"){
    s <- c(min(data[[variable]], na.rm=TRUE), 
           max(data[[variable]], na.rm=TRUE))
    if(variable %in% colnames(X)){
      X1[,variable] <- s[1]
      X2[,variable] <- s[2]
    }
    if(variable %in% colnames(Z)){
      Z1[,variable] <- s[1]
      Z2[,variable] <- s[2]
    }
  }  
  if(delta == "sd"){
    s <- sd(data[[variable]], na.rm=TRUE)
    if(variable %in% colnames(X)){
      X1[,variable] <- X1[,variable]- .5*s
      X2[,variable] <- X2[,variable]+ .5*s
    }
    if(variable %in% colnames(Z)){
      Z1[,variable] <- Z1[,variable]- .5*s
      Z2[,variable] <- Z2[,variable]+ .5*s
    }
  }
  if(delta == "unit"){
    s <- sd(data[[variable]], na.rm=TRUE)
    if(variable %in% colnames(X)){
      X1[,variable] <- X1[,variable]- .5
      X2[,variable] <- X2[,variable]+ .5
    }
    if(variable %in% colnames(Z)){
      Z1[,variable] <- Z1[,variable]- .5
      Z2[,variable] <- Z2[,variable]+ .5
    }
  }
  if(dist == "nb"){
    mu1 <- exp(X1 %*% t(as.matrix(B)))
    mu2 <- exp(X2 %*% t(as.matrix(B)))
    phi1 <- exp(Z1 %*% t(as.matrix(G)))
    phi2 <- exp(Z2 %*% t(as.matrix(G)))
    diff_mu <-  colMeans(mu2 - mu1)
    diff_sig <-  colMeans(sqrt(mu2 + mu2^2/phi2) - sqrt(mu1 + mu1^2/phi1))
  }
  if(dist == "normal"){
    mu1 <- X1 %*% t(as.matrix(B))
    mu2 <- X2 %*% t(as.matrix(B))
    sig1 <- exp(Z1 %*% t(as.matrix(G)))
    sig2 <- exp(Z2 %*% t(as.matrix(G)))
    diff_mu <-  colMeans(mu2 - mu1)
    diff_sig <-  colMeans(sig2-sig1)
  }
  
  pv_mu <- mean(diff_mu > 0)
  pv_sig <- mean(diff_sig > 0)
  pv_mu <- ifelse(pv_mu > .5, pv_mu, 1-pv_mu)
  pv_sig <- ifelse(pv_sig > .5, pv_sig, 1-pv_sig)
  
  lwr_lev <- (1-level)/2
  upr_lev <- 1-lwr_lev
  data.frame(
    vbl = c("mean", "sd"), 
    difference = c(mean(diff_mu), mean(diff_sig)), 
    lwr = c(unname(quantile(diff_mu, lwr_lev)), unname(quantile(diff_sig, lwr_lev))), 
    upr = c(unname(quantile(diff_mu, upr_lev)), unname(quantile(diff_sig, upr_lev))), 
    pval = c(pv_mu, pv_sig))
}

sigr <- function(obj, group, ...){
  as_draws_df(obj) %>% 
    select(contains(group)) %>% 
    select(-starts_with("r_")) %>% 
    summarise(across(everything(), {~ifelse(mean(.x >= 0) > .5, mean(.x >= 0), mean(.x <= 0))})) %>% 
    t()
}

sigf <- . %>% 
  as_draws_df(.) %>% 
  select(starts_with("b")) %>% 
  summarise(across(everything(), {~ifelse(mean(.x > 0) > .5, mean(.x > 0), mean(.x < 0))})) %>% 
  t()

fd.jags <- function(X, Z, b, g, variable, level=0.9, delta = c("sd", "range", "unit"), 
                    dist = c("nb", "normal")){
  delta <- match.arg(delta)
  dist <- match.arg(dist)
  B <- cbind(b, 1)
  G <- g
  X1 <- X2 <- X
  Z1 <- Z2 <- Z
  if(delta == "range"){
    s <- c(min(X[,variable], na.rm=TRUE), 
           max(X[,variable], na.rm=TRUE))
    if(variable %in% colnames(X)){
      X1[,variable] <- s[1]
      X2[,variable] <- s[2]
    }
    if(variable %in% colnames(Z)){
      Z1[,variable] <- s[1]
      Z2[,variable] <- s[2]
    }
  }  
  if(delta == "sd"){
    s <- NULL
    if(variable %in% colnames(X)){
      s <- sd(X[,variable], na.rm=TRUE)
    }
    if(variable %in% colnames(Z)){
      s <- sd(Z[,variable], na.rm=TRUE)
    }
    if(is.null(s))stop(paste0(variable, " not in either X or Z."))
    if(variable %in% colnames(X)){
      X1[,variable] <- X1[,variable]- .5*s
      X2[,variable] <- X2[,variable]+ .5*s
    }
    if(variable %in% colnames(Z)){
      Z1[,variable] <- Z1[,variable]- .5*s
      Z2[,variable] <- Z2[,variable]+ .5*s
    }
  }
  if(delta == "unit"){
    if(variable %in% colnames(X)){
      X1[,variable] <- X1[,variable]- .5
      X2[,variable] <- X2[,variable]+ .5
    }
    if(variable %in% colnames(Z)){
      Z1[,variable] <- Z1[,variable]- .5
      Z2[,variable] <- Z2[,variable]+ .5
    }
  }
  if(dist == "nb"){
    mu1 <- exp(X1 %*% t(as.matrix(B)))
    mu2 <- exp(X2 %*% t(as.matrix(B)))
    phi1 <- exp(Z1 %*% t(as.matrix(G)))
    phi2 <- exp(Z2 %*% t(as.matrix(G)))
    diff_mu <-  colMeans(mu2 - mu1)
    diff_sig <-  colMeans(sqrt(mu2 + mu2/phi2) - sqrt(mu1 + mu1/phi1))
  }
  if(dist == "normal"){
    mu1 <- X1 %*% t(as.matrix(B))
    mu2 <- X2 %*% t(as.matrix(B))
    sig1 <- exp(Z1 %*% t(as.matrix(G)))
    sig2 <- exp(Z2 %*% t(as.matrix(G)))
    diff_mu <-  colMeans(mu2 - mu1)
    diff_sig <-  colMeans(sig2-sig1)
  }
  
  pv_mu <- mean(diff_mu > 0)
  pv_sig <- mean(diff_sig > 0)
  pv_mu <- ifelse(pv_mu > .5, pv_mu, 1-pv_mu)
  pv_sig <- ifelse(pv_sig > .5, pv_sig, 1-pv_sig)
  
  lwr_lev <- (1-level)/2
  upr_lev <- 1-lwr_lev
  data.frame(
    vbl = c("mean", "sd"), 
    difference = c(mean(diff_mu), mean(diff_sig)), 
    lwr = c(unname(quantile(diff_mu, lwr_lev)), unname(quantile(diff_sig, lwr_lev))), 
    upr = c(unname(quantile(diff_mu, upr_lev)), unname(quantile(diff_sig, upr_lev))), 
    pval = c(pv_mu, pv_sig))
}

smry <- function(obj){
  pp <- prepare_predictions(obj)
  draws <- as_draws_df(obj)
  s <- summary(draws)
  s <- s %>% filter(grepl("^b_", variable))
  p <- sigf(obj)
  sigchar <- ifelse(p[,1] > .9, "*", "")
  entry <- s %>% mutate(sigchar = sigchar)  %>% transmute(variable = variable, 
                                                          upper = sprintf("%.2f%s", mean, sigchar), 
                                                          lower = sprintf("(%.2f, %.2f)",q5, q95)) %>% 
    tidyr::pivot_longer(-variable, names_to = "side", values_to = "col") %>% 
    dplyr::select(-side)
  entry
}
